# -*- coding: utf-8 -*-
# SimpleUserCF.py

from MovieLens import MovieLens
from surprise import KNNBasic
import heapq
from collections import defaultdict
from operator import itemgetter

testSubject = '85'
k = 10

# 加载数据集并计算用户相似度矩阵
ml = MovieLens()
data = ml.loadMovieLensLatestSmall()

trainSet = data.build_full_trainset()

sim_options = {
    'name': 'cosine',
    'user_based': True
}

# 模型是KNN
model = KNNBasic(sim_options=sim_options)
# 训练模型, 拟合数据
model.fit(trainSet)
# 计算相似度矩阵
simsMatrix = model.compute_similarities()


# 获取与测试用户最相似的前 N 个用户
# （另一种方法是选择相似度超过某个阈值的用户 - 尝试一下！）
testUserInnerID = trainSet.to_inner_uid(testSubject)
similarityRow = simsMatrix[testUserInnerID]  # 用户85的相似度矩阵

similarUsers = []  # 记录相似的用户
for innerID, score in enumerate(similarityRow):
    if (innerID != testUserInnerID):
        similarUsers.append((innerID, score))

# 找到前k个最相似的user
kNeighbors = heapq.nlargest(k, similarUsers, key=lambda t: t[1])

# 获取这些用户评分的项目，并根据用户相似度对每个项目的评分进行加权
candidates = defaultdict(float)
for similarUser in kNeighbors:
    innerID = similarUser[0]  # simi user id
    userSimilarityScore = similarUser[1]  # simi user score 和85号的相似分值
    theirRatings = trainSet.ur[innerID]  # simi user rating
    for rating in theirRatings:
        # 项目评分：相似用户评分 * 相似度分数
        candidates[rating[0]] += (rating[1] / 5.0) * userSimilarityScore

# 构建一个字典，记录该用户已经看过的项目
watched = {}
for itemID, rating in trainSet.ur[testUserInnerID]:
    watched[itemID] = 1  # 记录85号已经看过的movie, 没看过的为 movieID: 0

# 获取来自相似用户的最高评分项目：
pos = 0
# 根据 candidates候选字典的 [1](score) 排序, 然后遍历
for itemID, ratingSum in sorted(candidates.items(), key=itemgetter(1), reverse=True):
    # 85号没看过
    if not itemID in watched:
        movieID = trainSet.to_raw_iid(itemID)
        print(ml.getMovieName(int(movieID)), ratingSum)
        pos += 1
        if (pos > 10):
            break
